# ai_arm_controller.py
import json
from openai import OpenAI
from motor_action_3d import VirtualArm3D


class AIController:
    def __init__(self, api_key, base_url):
        """初始化AI控制器"""
        self.client = OpenAI(
            # 若没有配置环境变量，请用阿里云百炼API Key将下行替换为：api_key="sk-xxx",
            api_key=api_key,
            base_url=base_url,
        )

    def get_ai_response(self, user_input):
        """调用千问大模型获取指令"""
        prompt = """
        你是一个奥特曼机器人的控制系统。用户会发送指令，你需要理解并返回JSON格式的动作指令。
        可用的动作有: "hug"（拥抱）, "wave"（挥手）, "reset"（复位）。
        只返回JSON，不要其他内容。

        示例:
        输入: "抱一下", 输出: {"action": "hug"}
        输入: "挥挥手", 输出: {"action": "wave"}
        """
        try:
            response = self.client.chat.completions.create(
                model="qwen-plus",
                messages=[
                    {'role': 'system', 'content': prompt},
                    {'role': 'user', 'content': user_input}
                ],
                temperature=0.7,  # 控制随机性（0~1，越高越随机）
                max_tokens=100,  # 限制生成的最大长度
                top_p=0.9,  # 核采样（0~1，越高越多样）
            )

            return response.choices[0].message.content
        except Exception as e:
            print(f"调用大模型失败: {e}")
            return None


def main():
    # 配置
    api_key = "sk-a2ef3cb7fbaa4051af3128424d1d8f4c"  # 替换为您的API Key
    base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
    # 初始化
    ai_controller = AIController(api_key=api_key, base_url=base_url)
    virtual_arm = VirtualArm3D()

    print("=== AI奥特曼控制系统已启动 ===")
    print("输入指令示例: '拥抱一下', '挥挥手', '复位'")
    print("输入 'exit' 退出程序\n")

    try:
        while True:
            # 获取用户输入
            user_input = input("请输入指令: ").strip()
            if user_input.lower() == 'exit':
                break
            if not user_input:
                continue
            print(f"🧠 分析指令: {user_input}")
            # 调用大模型
            ai_response = ai_controller.get_ai_response(user_input)
            if not ai_response:
                print("❌ 大模型无响应")
                continue
            print(f"📨 大模型回复: {ai_response}")
            # 解析JSON指令
            try:
                command = json.loads(ai_response)
                action = command.get("action")

                if action in ["hug", "wave", "reset"]:
                    print(f"🎯 执行动作: {action}")
                    virtual_arm.perform_action(action)
                else:
                    print(f"⚠️ 未知动作: {action}")

            except json.JSONDecodeError:
                print("❌ 大模型返回了非JSON格式的内容")

    except KeyboardInterrupt:
        print("\n程序被用户中断")
    finally:
        virtual_arm._reset()
        print("=== 系统已关闭 ===")


if __name__ == "__main__":
    main()